from .enum_class import LLMType, EmbeddingType
from .llm import LLMHandler
from .embedding import EmbeddingHandler
from dotenv import load_dotenv
import os
from llama_index.core import Settings
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core import StorageContext, load_index_from_storage
from llama_index.core.tools import QueryEngineTool
from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector, PydanticSingleSelector


class MainHandler:
    """
    主处理器，用于整合所有功能
    """

    def __init__(
        self,
        llm_type: LLMType,
        embedding_type: EmbeddingType,
        **kwargs,
    ):
        """
        初始化主处理器
        参数:
            llm_type: LLM类型，目前支持的可选值为 Local_HuggingFace 或 Online_OpenRouter 或 Online_OpenAILike
            embedding_type: Embedding类型，目前支持的可选值为 Local_HuggingFace
            **kwargs: 其他参数，用于传递给MainHandler的初始化参数
            可选参数举例:
                config_path: 配置文件路径，默认".env"
                context_window: 上下文窗口大小，默认4096
                num_output_cache: 输出缓冲区大小，默认512
                llm_model_name: LLM模型名称，默认"Langboat/bloom-1b4-zh"
                embedding_model_name: Embedding模型名称，默认"BAAI/bge-small-zh-v1.5"
                以及其他模型参数等
        """
        # 加载环境变量
        if "config_path" in kwargs:
            config_path = kwargs.get("config_path")
            kwargs.pop("config_path")
            load_dotenv(config_path)
        else:
            load_dotenv()

        # 是否支持函数调用（Function Calling）
        self.is_function_calling_model = kwargs.get("is_function_calling_model")

        # 设置上下文窗口大小
        Settings.context_window = kwargs.get("context_window", 4096)
        if "context_window" in kwargs:
            kwargs.pop("context_window")
        # 设置输出缓冲区大小
        Settings.num_output_cache = kwargs.get("num_output_cache", 512)
        if "num_output_cache" in kwargs:
            kwargs.pop("num_output_cache")

        # 初始化类的 llm_type
        self.llm_type = llm_type

        # 配置全局LLM模型
        llm = self.get_llm(llm_type, **kwargs)
        Settings.llm = llm
        if self.llm_type == LLMType.Local_HuggingFace:
            print(f"LLM模型加载完成: {llm.model_name}")
        else:
            print(f"LLM模型加载完成: {llm}")

        # 配置全局Embedding模型
        embedding = self.get_embedding(embedding_type, **kwargs)
        Settings.embed_model = embedding
        print(f"Embedding模型加载完成: {embedding}")

    def get_config(self, key: str):
        """
        获取配置
        """
        return os.getenv(key)

    def get_openailike_api_key(self):
        """
        获取兼容OpenAI格式模型的API密钥
        优先级:
        1. DEEPSEEK_API_KEY
        2. VOLCENGINE_API_KEY
        3. SILICONFLOW_API_KEY
        """
        # 按优先级顺序检查各个API key
        api_keys = [
            "DEEPSEEK_API_KEY",
            "VOLCENGINE_API_KEY",
            "SILICONFLOW_API_KEY",
        ]

        for key in api_keys:
            api_key = self.get_config(key)
            if api_key:
                return api_key

        return None

    def get_llm(self, llm_type: LLMType, **kwargs):
        """
        获取封装好的LLM模型
        参数:
            llm_type: LLM类型，目前支持的可选值为 Local_HuggingFace 或 Online_OpenRouter 或 Online_OpenAILike
            **kwargs: 其他参数，用于传递给LLM模型的初始化参数
            可选参数举例:
                llm_model_name: 模型名称，如果不传则使用默认值
        """
        # 处理所有的api_key
        match llm_type:
            case LLMType.Online_OpenRouter:
                openrouter_api_key = self.get_config("OPENROUTER_API_KEY")
                if openrouter_api_key:
                    kwargs["api_key"] = openrouter_api_key
                else:
                    raise ValueError("OPENROUTER_API_KEY 未设置")
            case LLMType.Online_OpenAILike:
                openailike_api_key = self.get_openailike_api_key()
                if openailike_api_key:
                    kwargs["api_key"] = openailike_api_key
                else:
                    raise ValueError(
                        "API_KEY 未设置，目前支持 DEEPSEEK_API_KEY、VOLCENGINE_API_KEY、SILICONFLOW_API_KEY"
                    )
            case _:
                pass

        # 创建LLMHandler实例
        llm_handler = LLMHandler(llm_type, **kwargs)
        # 返回封装好的LLM模型
        return llm_handler.get_llm()

    def get_embedding(self, embedding_type: EmbeddingType, **kwargs):
        """
        获取封装好的Embedding模型
        参数:
            embedding_type: Embedding类型，目前支持的可选值为 Local_HuggingFace
            **kwargs: 其他参数，用于传递给Embedding模型的初始化参数
            可选参数:
                embedding_model_name: 模型名称，如果不传则使用默认值
        """
        # 创建EmbeddingHandler实例
        embedding_handler = EmbeddingHandler(embedding_type, **kwargs)
        # 返回封装好的Embedding模型
        return embedding_handler.get_embedding()

    def create_vector_database(
        self,
        knowledge_path: str,
        recursive: bool = True,
        persist_dir: str = "./storage",
        show_progress: bool = True,
    ):
        """
        创建向量数据库
        参数:
            knowledge_path: 知识库路径
            recursive: 是否递归加载文档，默认True，文件夹下的子孙文件都会被加载
            show_progress: 创建索引时是否显示进度，默认True
            persist_dir: 向量数据持久化存储路径，默认"./storage"
        """
        # 加载文档
        documents = SimpleDirectoryReader(
            knowledge_path, recursive=recursive
        ).load_data()

        # 创建索引
        index = VectorStoreIndex.from_documents(documents, show_progress=show_progress)

        # 持久化存储
        index.storage_context.persist(persist_dir=persist_dir)

    def get_vector_query_engine_tools(
        self,
        vector_database_path: str = "./storage",
        description: str = "这是查询引擎的示例描述，用于回答关于知识库的问题",
        streaming: bool = True,
    ):
        """
        获取向量查询引擎工具
        参数:
            vector_database_path: 向量数据持久化存储路径，默认"./storage"
            description: 查询引擎的描述，默认"这是查询引擎的示例描述，用于回答关于知识库的问题"
            streaming: 是否启用流式输出，默认True
        返回:
            vector_query_engine_tool: 向量查询引擎工具
        """
        # 判断目录是否存在
        if not os.path.exists(vector_database_path):
            raise FileNotFoundError(f"向量数据库路径不存在：{vector_database_path}")

        # 加载索引
        storage_context = StorageContext.from_defaults(persist_dir=vector_database_path)
        index = load_index_from_storage(storage_context)

        # 创建查询引擎工具
        vector_query_engine_tool = QueryEngineTool.from_defaults(
            query_engine=index.as_query_engine(
                similarity_top_k=3,  # 相似度最高的3个结果
                response_mode="tree_summarize",  # 结果聚合模式：树形聚合
                streaming=streaming,  # 是否启用流式输出
            ),
            description=description,
        )

        return vector_query_engine_tool

    def get_router_query_engine(
        self,
        vector_database_path: str = "./storage",
        description: str = "这是查询引擎的示例描述，用于回答关于知识库的问题",
        is_function_calling_model: bool = False,
        streaming: bool = True,
    ):
        """
        获取路由查询引擎
        参数:
            vector_database_path: 向量数据持久化存储路径，默认"./storage"
            description: 查询引擎的描述，默认"这是查询引擎的示例描述，用于回答关于知识库的问题"
            is_function_calling_model: 是否为支持函数调用的模型，默认False
            streaming: 是否启用流式输出，默认True
        返回:
            RouterQueryEngine: 路由查询引擎实例
        """
        # 判断目录是否存在
        if not os.path.exists(vector_database_path):
            raise FileNotFoundError(f"向量数据库路径不存在：{vector_database_path}")

        # 加载向量数据库引擎
        vector_query_engine_tool = self.get_vector_query_engine_tools(
            vector_database_path, description, streaming=streaming
        )

        is_function_calling_model = (
            self.is_function_calling_model
            if self.is_function_calling_model
            else is_function_calling_model
        )
        if is_function_calling_model:
            # 对于支持函数调用的模型，使用PydanticSingleSelector，通过函数调用（Function Calling）选择下一步最合适调用的工具
            query_engine = RouterQueryEngine(
                selector=PydanticSingleSelector.from_defaults(),
                query_engine_tools=(vector_query_engine_tool,),
            )

        else:
            # 对于不支持函数调用（Function Calling）的模型，使用LLMSingleSelector，通过大语言模型选择下一步最合适调用的工具
            # 因此我们需要通过工具调用的返回的提示词来明确要求JSON格式输出，但也有一定概率因为模型连接等原因导致报错，建议在调用层做异常处理

            # 自定义选择器提示模板，明确要求JSON格式输出
            custom_selector_prompt_template = (
                "Some choices are given below. It is provided in a numbered list "
                "(1 to {num_choices}), "
                "where each item in the list corresponds to a summary.\n"
                "---------------------\n"
                "{context_list}"
                "\n---------------------\n"
                "Using only the choices above and not prior knowledge, return "
                "the choice that is most relevant to the question: '{query_str}'\n"
                "IMPORTANT: Your response must be a valid JSON object in the exact format shown below. "
                "Do not include any explanations, only provide the JSON.\n"
                'Format: {"selections": [{"choice": <number>, "reason": "<explanation>"}]}\n'
                "where <number> is the number of the chosen option (1 to {num_choices}) and "
                "<explanation> is a brief reason for this choice.\n"
            )

            # 创建路由查询引擎
            query_engine = RouterQueryEngine(
                selector=LLMSingleSelector.from_defaults(
                    prompt_template_str=custom_selector_prompt_template
                ),
                query_engine_tools=(vector_query_engine_tool,),
            )

        return query_engine
